import tkinter as tk
from tkinter import ttk
from tkinter import messagebox
from tkinter.filedialog import askdirectory, asksaveasfilename, askopenfilename
from PyPDF2 import PdfMerger, PdfReader, PdfWriter
from pathlib import Path
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class PDFMergerApp:
    def __init__(self, root):
        self.root = root
        self.root.title("PDF文件合并/拆分")

        # 合并相关变量
        self.merge_dirpath = tk.StringVar()
        self.merge_output = tk.StringVar()

        # 拆分相关变量
        self.split_filepath = tk.StringVar()
        self.split_outputdir = tk.StringVar()
        self.split_pages = tk.IntVar(value=1)

        self.create_widgets()

    def create_widgets(self):
        # 合并功能区
        lfr1 = ttk.LabelFrame(self.root, text='合并')
        lfr1.grid(row=0, column=0, sticky='nsew', padx=10, pady=5)

        # 选择合并目录
        ttk.Label(lfr1, text='选择目录:').grid(row=0, column=0, padx=5, pady=5)
        ttk.Entry(lfr1, textvariable=self.merge_dirpath, width=40).grid(row=0, column=1, padx=5, pady=5)
        ttk.Button(lfr1, text='浏览', command=self.select_merge_dir).grid(row=0, column=2, padx=5, pady=5)

        # 合并输出文件
        ttk.Label(lfr1, text='保存文件:').grid(row=1, column=0, padx=5, pady=5)
        ttk.Entry(lfr1, textvariable=self.merge_output, width=40).grid(row=1, column=1, padx=5, pady=5)
        ttk.Button(lfr1, text='合并', command=self.merge_files).grid(row=1, column=2, padx=5, pady=5)

        # 拆分功能区
        lfr2 = ttk.LabelFrame(self.root, text='拆分')
        lfr2.grid(row=1, column=0, sticky='nsew', padx=10, pady=5)

        # 选择拆分文件
        ttk.Label(lfr2, text='选择文件:').grid(row=0, column=0, padx=5, pady=5)
        ttk.Entry(lfr2, textvariable=self.split_filepath, width=40).grid(row=0, column=1, padx=5, pady=5)
        ttk.Button(lfr2, text='浏览', command=self.select_split_file).grid(row=0, column=2, padx=5, pady=5)

        # 拆分页数设置
        ttk.Label(lfr2, text='每份页数:').grid(row=1, column=0, padx=5, pady=5)
        ttk.Spinbox(lfr2, from_=1, to=100, textvariable=self.split_pages, width=5).grid(row=1, column=1, sticky='w', padx=5, pady=5)

        # 选择输出目录
        ttk.Label(lfr2, text='输出目录:').grid(row=2, column=0, padx=5, pady=5)
        ttk.Entry(lfr2, textvariable=self.split_outputdir, width=40).grid(row=2, column=1, padx=5, pady=5)
        ttk.Button(lfr2, text='浏览', command=self.select_split_dir).grid(row=2, column=2, padx=5, pady=5)
        ttk.Button(lfr2, text='拆分', command=self.split_files).grid(row=3, column=1, pady=5)

        # 配置网格布局
        self.root.columnconfigure(0, weight=1)
        lfr1.columnconfigure(1, weight=1)
        lfr2.columnconfigure(1, weight=1)

    def select_merge_dir(self):
        directory = askdirectory()
        if directory:
            self.merge_dirpath.set(directory)

    def select_split_file(self):
        filepath = askopenfilename(filetypes=[("PDF文件", "*.pdf")])
        if filepath:
            self.split_filepath.set(filepath)

    def select_split_dir(self):
        directory = askdirectory()
        if directory:
            self.split_outputdir.set(directory)

    def merge_files(self):
        try:
            src_dir = Path(self.merge_dirpath.get())
            if not src_dir.is_dir():
                raise ValueError("请选择有效的合并目录")

            output_path = asksaveasfilename(
                defaultextension=".pdf",
                filetypes=[("PDF文件", "*.pdf")]
            )
            if not output_path:
                return

            merger = PdfMerger()
            for pdf_file in sorted(src_dir.glob("*.pdf")):
                merger.append(pdf_file)

            with open(output_path, "wb") as f:
                merger.write(f)

            self.merge_output.set(output_path)
            messagebox.showinfo("成功", f"文件已合并保存至:\n{output_path}")
        except Exception as e:
            messagebox.showerror("错误", str(e))

    def split_files(self):
        try:
            input_path = Path(self.split_filepath.get())
            output_dir = Path(self.split_outputdir.get())
            pages = self.split_pages.get()

            if not input_path.is_file():
                raise ValueError("请选择有效的PDF文件")
            if not output_dir.exists():
                output_dir.mkdir(parents=True)
            if pages < 1:
                raise ValueError("每份页数必须大于0")

            with open(input_path, "rb") as f:
                reader = PdfReader(f)
                total_pages = len(reader.pages)

                # 计算需要的数字位数（动态前导零）
                max_split = (total_pages + pages - 1) // pages
                zero_padding = len(str(max_split))

                for i in range(0, total_pages, pages):
                    writer = PdfWriter()
                    end = min(i + pages, total_pages)
                    for j in range(i, end):
                        writer.add_page(reader.pages[j])

                    # 使用动态前导零格式
                    output_path = output_dir / f"page_{(i//pages + 1):0{zero_padding}d}.pdf"
                    with open(output_path, "wb") as out_f:
                        writer.write(out_f)

            messagebox.showinfo("成功", f"文件已拆分保存至:\n{output_dir}")
        except Exception as e:
            messagebox.showerror("错误", str(e))

if __name__ == "__main__":
    root = tk.Tk()
    root.geometry("450x400")
    app = PDFMergerApp(root)
    root.mainloop()